import abc
import pandas as pd
import docx
import PyPDF2
import markdown
import json
from bs4 import BeautifulSoup
from typing import Union

# 假设 utils.Mydb 模块中的 db 已经正确初始化
from utils.Mydb import db


# 定义抽象加载器基类
class AbstractLoader(abc.ABC):
    @abc.abstractmethod
    def load(self, file_path):
        pass


# 具体的 Markdown 加载器
class MarkdownLoader(AbstractLoader):
    def load(self, file_path):
        with open(file_path, 'r', encoding='utf - 8') as f:
            text = f.read()
            html = markdown.markdown(text)
            return html

# 具体的TXT加载器
class TXTLoader(AbstractLoader):
    def load(self, file_path):
        with open(file_path, 'r', encoding='utf - 8') as f:
            return f.read()


# 具体的 CSV 加载器
class CSVLoader(AbstractLoader):
    def load(self, file_path):
        data = pd.read_csv(file_path)
        return data.to_dict('records')


# 具体的 HTML 加载器
class HTMLLoader(AbstractLoader):
    def load(self, file_path):
        with open(file_path, 'r', encoding='utf - 8') as f:
            soup = BeautifulSoup(f, 'lxml')
            return soup.get_text()


# 具体的 JSON 加载器
class JSONLoader(AbstractLoader):
    def load(self, file_path):
        with open(file_path, 'r', encoding='utf - 8') as f:
            data = json.load(f)
            return data


# 具体的 PDF 加载器
class PDFLoader(AbstractLoader):
    def load(self, file_path):
        pdf_file_obj = open(file_path, 'rb')
        pdf_reader = PyPDF2.PdfReader(pdf_file_obj)
        text = ""
        for page_num in range(len(pdf_reader.pages)):
            page_obj = pdf_reader.pages[page_num]
            text += page_obj.extract_text()
        pdf_file_obj.close()
        return text


# 具体的 Docx 加载器
class DocxLoader(AbstractLoader):
    def load(self, file_path):
        doc = docx.Document(file_path)
        content = ""
        for para in doc.paragraphs:
            content += para.text + "\n"
        return content


# 定义抽象工厂基类
class AbstractLoaderFactory(abc.ABC):
    @abc.abstractmethod
    def create_loader(self, file_path: str) -> Union[AbstractLoader, None]:
        pass


# 具体的文件加载器工厂类
class FileLoaderFactory(AbstractLoaderFactory):
    def create_loader(self, file_path: str) -> Union[AbstractLoader, None]:
        file_extension = file_path.split('.')[-1].lower()
        if file_extension =='md':
            return MarkdownLoader()
        elif file_extension == 'csv':
            return CSVLoader()
        elif file_extension == 'html':
            return HTMLLoader()
        elif file_extension == 'json':
            return JSONLoader()
        elif file_extension == 'pdf':
            return PDFLoader()
        elif file_extension in ['doc', 'docx']:
            return DocxLoader()
        elif file_extension == 'txt':  # 新增判断
            return TXTLoader()
        return None

    @staticmethod
    def load_to_vector_db(file_path):
        factory = FileLoaderFactory()
        loader = factory.create_loader(file_path)
        if loader:
            data = loader.load(file_path)
            if isinstance(data, list):
                data = " ".join(str(item) for item in data)
            elif isinstance(data, dict):
                data = str(data)
            with open('./utils/temp.txt', 'a', encoding='utf - 8') as f:
                f.write(data)
            return f"Successfully loaded {file_path} to vector database."
        else:
            return f"Unsupported file format: {file_path}"


# # 正确的调用方式
# file_path = r"C:\Users\meng\Desktop\python前后端分离配置.txt"
# result = FileLoaderFactory.load_to_vector_db(file_path)
# print(result)